# Turning other markets sold data into ind. lvl. adj. mat. ---------------------
make_ind_adj <- function(df, ids){
  
  othr_mkts <- data.frame(t(sapply(df$`Other Markets in Our Sample`, 
                                   function(x){
                                     mkts <- str_split(x, ", ")[[1]]
                                     
                                     1*(market_names %in% mkts)
                                     
                                   })))
  
  row.names(othr_mkts) <- NULL
  
  colnames(othr_mkts) <- market_names
  
  df <- cbind.data.frame(df, othr_mkts)
  
  dat <- merge(df, ids, all = T)
  
  dat <- dat[, -c(2:6)]
  
  dat[is.na(dat)] <- 0
  
  dat <- dat[order(dat$market), ]
  
  dat
}

# Miguel Spillover Functions ---------------------------------------------------
#functions to calculate within treatment group vendor totals for each market
#at a certain distance from the market

#function to get adjacency matrix
get_adj <- function(dist_mat, d){
  1*(dist_mat <= d & dist_mat > 0)
}

get_d_total <- function(adj_mat, d, gr_var, var_tbs, var_name,
                        incl_slf = T){
  #browser()
  grs <- unique(gr_var)
  
  sums <- lapply(grs, function(x){
    adj_mat_gr <- apply(adj_mat, 1, function(y) 1*(gr_var == x)*y)
    
    adj_mat_gr %*% var_tbs
    
  })
  
  if(incl_slf) diag(adj_mat) <- 1
  
  tot <- adj_mat %*% var_tbs
  
  sums <- data.frame(sums, tot)
  
  sums_names <- sapply(c(as.character(grs), "all"), function(x){
    paste0(var_name, "_", d, "_", x)
  })
  
  names(sums) <- sums_names
  
  sums
}

make_adjs <- function(dist_mat, d, rings = T){
  
  adjs <- lapply(d, function(x){
    get_adj(dist_mat, x)
  })
  
  if(rings & !(length(d) == 1)){
    #getting differences in adjacency matrices to get "rings"
    for(i in 2:length(adjs)){
      adjs[[i]] <- adjs[[i]] - adjs[[i - 1]]
      
    }
  }
  
  adjs
}


get_ring_tot <- function(dist_mat, d, gr_var, var_tbs, 
                         var_name,
                         incl_slf = T){
  
  adjs <- make_adjs(dist_mat, d)
  
  map2_dfc(adjs, d, get_d_total, gr_var = gr_var, 
           var_name = var_name, var_tbs = var_tbs, incl_slf = incl_slf)
  
}


get_ring_tots <- function(dist_mat, d, gr_var, vars_tbs, id,
                          incl_slf = T){
  
  dfs <- lapply(1:length(vars_tbs), function(x){
    df <- get_ring_tot(dist_mat, d, gr_var,
                       vars_tbs[[x]], names(vars_tbs)[x])
  })
  
  dfs <- do.call(cbind, dfs)
  
  data.frame(id = id, 
             group = gr_var,
             dfs)
}

# IPW Spillover Functions ------------------------------------------------------
### IPW FUNCTIONS

#function to simulate block random assignment
block_ra <- function(blocks, N_within, treat_m, as_matrix = T, ind_lvl = F,
                     n_per_unit = NULL){
  #browser()
  assigned <- rep(NA, length(blocks))
  
  for(i in unique(blocks)){
    assigned[blocks == i] <-  sample(1:treat_m, N_within)
  }
  
  if(ind_lvl){
    if(is.null(n_per_unit) | is.character(n_per_unit)){
      stop("n_per_unit is missing")
    }
    
    if(length(n_per_unit) == 1){
      n_per_unit <- rep(n_per_unit, length(blocks))
    }
    
    assigned <- rep(assigned, times = n_per_unit)
  }
  
  if(as_matrix){
    #getting 1/0 vectors
    #assumes 1 = control
    assigned <- sapply(2:treat_m, function(x){
      1*(assigned == x)
    })    
  }
  
  assigned
}

#function to get true condition, based on adjacency matrix 
get_condition <- function(assigned, adjmat, subs_id = 1:nrow(adjmat),
                          ind_lvl = F, n_per_unit = NULL){
  #browser()
  
  assigned <- assigned[subs_id,]
  
  exposure <- adjmat %*% assigned
  
  #condition <- rep("0000", nrow(adjmat))
  
  exposed <- 1*(exposure > 0)
  
  if(ind_lvl){
    if(length(n_per_unit) == 1){
      n_per_unit <- rep(n_per_unit, length(n_per_unit))
    }
    
    assigned <- as.data.frame(assigned)[rep(1:nrow(assigned),
                                            times = n_per_unit[subs_id]),]
    assigned <- as.matrix(assigned)
  }
  
  #to make a buffer row
  assigned <- cbind(assigned, rep("_", nrow(assigned)))
  
  apply(cbind(assigned, exposed), 1, paste0, sep = "", collapse = "")
}

#function that repeats block_ra sim times and stores output in a list
block_ra_rep <- function(blocks, N_within, treat_m, sim = 1000, as_matrix = T,
                         verbose = F, ind_lvl = F, n_per_unit = NULL){
  
  lapply(1:sim, function(x){
    
    if(verbose){
      cat("Working on simulation:", x, "\n")
    }
    
    block_ra(blocks = blocks, N_within = N_within,
             treat_m = treat_m, as_matrix = as_matrix, ind_lvl = ind_lvl, 
             n_per_unit = n_per_unit)
    
  })
}

#get_condition but for a list of assigneds
get_conditions <- function(assigneds, adjmat, subs_id = 1:nrow(adjmat),
                           ind_lvl = F, n_per_unit = NULL){
  
  sapply(assigneds, get_condition, adjmat = adjmat, subs_id = subs_id,
         ind_lvl = ind_lvl, n_per_unit = n_per_unit)
  
}

#function to calculate probabilities of being in each treatment condition,
#based on simulations
get_sim_probs <- function(cond_mat, conds_pos){
  #browser()
  prob_df <- data.frame(sapply(conds_pos, function(x){
    #browser()
    rowMeans(cond_mat == x)
  }))
  
  names(prob_df) <- conds_pos
  
  prob_df
}

#function to get return one vector that reflects probability of being in unit it was assigned
get_prob <- function(prob_df, cond_vec, conds_pos){
  prob <- rep(NA, nrow(prob_df))
  #browser()
  for(i in conds_pos){
    prob[cond_vec == i] <- as.data.frame(prob_df)[cond_vec == i, 
                                                  grepl(i, names(prob_df))]
  }
  
  prob
}

#function that combines all of the above

get_spillIPW <- function(conds, distances, dist_mat, treat_assigned,
                         subset_unit = 1:length(treat_assigned), blocks, N_within,
                         treat_m, sim, ind_lvl = F, 
                         n_per_unit = NULL, ind_adj, verbose = F,
                         id_vec){
  #browser()
  subset <- subset_unit
  
  if(ind_lvl) {
    if(length(n_per_unit) == 1){
      n_per_unit <- rep(n_per_unit, length(blocks))
    }
    
    subset <- rep(subset_unit, times = n_per_unit)
  }
  
  
  dist_str <- sapply(distances, function(x){
    paste0("d",x)
  })
  
  #adjacency matrices for distances
  adj_ipw <- lapply(distances, function(x){
    get_adj(dist_mat, x)
  })
  
  if(ind_lvl){
    adj_ipw <- lapply(adj_ipw, function(x){
      #browser()
      x <- data.frame(x)
      x <- x[rep(1:(length(blocks[subset_unit])),
                 times = n_per_unit[subset_unit]), ]
      x <- as.matrix(x)
      x <- x + ind_adj[subset, subset_unit]
      1*(x > 0)
      
    })
    
  }
  
  #the true condition for all units for each set of distances
  true_cond_IPW <- lapply(adj_ipw, function(x) {
    get_condition(treat_assigned, x, 
                  subs_id = subset_unit,
                  ind_lvl = ind_lvl,
                  n_per_unit =n_per_unit)
  })
  
  #the probability of being in each of the 32 conditions for each market,
  #based on simulation
  cond_probs <- lapply(adj_ipw, function(.x){
    block_ra_rep(blocks, N_within, treat_m, sim, ind_lvl = F,
                 n_per_unit = n_per_unit, verbose = verbose) %>% 
      get_conditions(.x, subs_id = subset_unit,
                     ind_lvl = ind_lvl, n_per_unit = n_per_unit) %>%
      get_sim_probs(conds)
  }) 
  
  cond_probs <- map2(cond_probs, dist_str, function(.x, .y){
    names(.x) <- sapply(names(.x), function(t){
      paste(.y, t, sep = "_")
    })
    .x
  })
  
  #the probability of being in the observed condition (will be used as weights)
  probs_spill <-  map2(cond_probs, true_cond_IPW, function(.x, .y){
    get_prob(.x, .y, conds)
  }) 
  
  #data frame that combines true conditions and probability of being in condition
  spills_IPW <- data.frame(as.data.frame(true_cond_IPW), 
                           as.data.frame(probs_spill))
  
  names(spills_IPW) <- apply(cbind(rep(c("cond", 
                                         "prob"), 
                                       each = length(dist_str)),
                                   rep(dist_str, times = 2)), 1,
                             paste, collapse = "_", sep = "")
  
  #adding in market name
  spills_IPW$id <- id_vec[subset]
  spills_IPW <- spills_IPW[, c("id", names(spills_IPW)[-ncol(spills_IPW)])]
  
  
  #now putting everything together (name, true conditions, true condition probs,
  #all condition probs)
  cond_probs[[(length(cond_probs) + 1)]] <- spills_IPW
  
  do.call(cbind, cond_probs)
}

#function to get mode of a vector
getmode <- function(v, na.rm = F) {
  if(na.rm) v <- v[!is.na(v)]
  uniqv <- unique(v)
  uniqv[which.max(tabulate(match(v, uniqv)))]
}
